struct DecisionTreeConfig
{
    DecisionTreeConfig() {}
    DecisionTreeConfig(UINT _maxTreeDepth, UINT _leafNodeCountCutoff, UINT _maxTestsPerDimension)
    {
        maxTreeDepth = _maxTreeDepth;
        leafNodeCountCutoff = _leafNodeCountCutoff;
        maxTestsPerDimension = _maxTestsPerDimension;
    }

    UINT maxTreeDepth;
    UINT leafNodeCountCutoff;
    UINT maxTestsPerDimension;
};

class DecisionTreeNode
{
public:
    DecisionTreeNode()
    {
        for(UINT childIndex = 0; childIndex < childCount; childIndex++)
        {
            _children[childIndex] = NULL;
        }
    }

    ~DecisionTreeNode()
    {
        for(UINT childIndex = 0; childIndex < childCount; childIndex++)
        {
            if(_children[childIndex] != NULL)
            {
                delete _children[childIndex];
                _children[childIndex] = NULL;
            }
        }
    }

    void Evaluate(const LearnerInput &input, double &probabilityClass0) const
    {
        if(Leaf())
        {
            probabilityClass0 = double(_class0Count) / double(_totalCount);
        }
        else
        {
            return _children[ComputeChildIndex(input)]->Evaluate(input, probabilityClass0);
        }
    }

    void Train(const DecisionTreeConfig &config, const LearnerDataset &examples, const vector<UINT> &activeExampleIndices, UINT depthRemaining);

private:
    inline UINT ComputeChildIndex(const LearnerInput &input) const
    {
        if(input.data[_decisionVariableIndex] <= _decisionVariableValue)
        {
            return 0;
        }
        else
        {
            return 1;
        }
    }

    inline bool Leaf() const
    {
        return (_children[0] == NULL);
    }

    static inline double Entropy(double p)
    {
        if(p <= 0.0 || p >= 1.0)
        {
            return 0.0;
        }
        else
        {
            double np = 1.0 - p;
            return -(p * log(p) + np * log(np));
        }
    }

    bool ChooseVariableAndThreshold(const DecisionTreeConfig &config, const LearnerDataset &examples, const vector<UINT> &activeExampleIndices);
    void InitLeafProbabilities(const DecisionTreeConfig &config, const LearnerDataset &examples, const vector<UINT> &activeExampleIndices);
    void TrainChild(const DecisionTreeConfig &config, const LearnerDataset &examples, const vector<UINT> &activeExampleIndices, UINT childIndex, UINT DepthRemaining);
    
    friend OutputDataStream& operator << (OutputDataStream &stream, const DecisionTreeNode &node);
    friend InputDataStream& operator >> (InputDataStream &stream, DecisionTreeNode &node);

    UINT _decisionVariableIndex;
    int _decisionVariableValue;

    static const UINT childCount = 2;
    DecisionTreeNode *_children[childCount];

    //
    // Leaf nodes
    //
    UINT _class0Count;
    UINT _totalCount;
};

class DecisionTree : public BinaryClassifier
{
public:
    DecisionTree() {}
    DecisionTree(const DecisionTreeConfig &config)
    {
        _config = config;
    }
    void Configure(const DecisionTreeConfig &config)
    {
        _config = config;
    }

    void Train(const LearnerDataset &dataset);

    double EvaluateProbability(const LearnerInput &input) const
    {
        double probabilityClass0 = 0.0;
        _root.Evaluate(input, probabilityClass0);
        return probabilityClass0;
    }

    int Evaluate(const LearnerInput &input) const
    {
        double probabilityClass0 = 0.0;
        _root.Evaluate(input, probabilityClass0);
        if(probabilityClass0 >= 0.5) return 0;
        return 1;
    }

private:
    friend OutputDataStream& operator << (OutputDataStream &stream, const DecisionTree &tree);
    friend InputDataStream& operator >> (InputDataStream &stream, DecisionTree &tree);

    DecisionTreeConfig _config;

    DecisionTreeNode _root;
};

OutputDataStream& operator << (OutputDataStream &stream, const DecisionTreeNode &node);
InputDataStream& operator >> (InputDataStream &stream, DecisionTreeNode &node);

OutputDataStream& operator << (OutputDataStream &stream, const DecisionTree &tree);
InputDataStream& operator >> (InputDataStream &stream, DecisionTree &tree);
